# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import math
import sys

from coin.base.datetime_util import to_datetime, to_timestamp_int
from coin.exchange.bitfinex_v2.kr_rest.product import BitfinexProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.calculator import MovingSum


def ceil_price(price, tick_size):
  return math.ceil(price / tick_size) * tick_size


def floor_price(price, tick_size):
  return math.floor(price / tick_size) * tick_size


class Feature(object):
  def __init__(self,
               product_1,
               product_2,
               fee_1,
               fee_2,
               edge,
               basis_ma_window,
               lot_size,
               stack,
               start_ts=0,
               logger=None):
    self._logger = logger or logging.getLogger('ArbMM')

    # Parameters
    # product_1: Aggression
    # product_2: Passive
    self._product_1 = product_1
    self._product_2 = product_2

    self._fee_1 = fee_1
    self._fee_2 = fee_2

    self._basis_ma_window = basis_ma_window
    self._edge = edge
    self._stack_2 = stack
    self._lot_size_2 = float(lot_size)

    self._logger.info('window=%dmin edge=%fbp stack=%d lot_size=%f',
                      self._basis_ma_window / (60 * (10**9)),
                      self._edge * 10000,
                      self._stack_2,
                      self._lot_size_2)

    self._tick_size_2 = 1.0
    self._true_book_qty_1 = 1000
    self._true_book_qty_2 = 1000

    self._start_ts = to_timestamp_int(start_ts)
    self._last_fire_ts = 0

    self._book_1 = None
    self._book_2 = None

    # Basis calculator
    self._ask_basis_ma = MovingSum(self._basis_ma_window)
    self._bid_basis_ma = MovingSum(self._basis_ma_window)

    self._pos = 0
    self._pnl = 0
    self._min_pnl = 0
    self._max_pnl = 0
    self._num_trade = 0

    print('ts,human_ts,ask0_1,bid0_1,askt_1,bidt_1,'
          'ask0_2,bid0_2,askt_2,bidt_2,'
          'ask_basis_ma,bid_basis_ma,'
          'pos,pnl')

  def on_book_reset(self, book_builder_name, book_builder):
    if book_builder_name == get_subscription_key(self._product_1):
      book_builder.subscribe(self._product_1, self.on_book_1)
    elif book_builder_name == get_subscription_key(self._product_2):
      book_builder.subscribe(self._product_2, self.on_book_2)
    else:
      raise ValueError(book_builder_name)

  def on_book_1(self, book):
    self._book_1 = book
    self.calculate()

  def on_book_2(self, book):
    self._book_2 = book
    self.calculate()

  def calculate(self):
    if self._book_1 is None or self._book_2 is None:
      return

    book_ts = max(self._book_1.timestamp, self._book_2.timestamp)
    min_ts = min(self._book_1.timestamp, self._book_2.timestamp)
    if book_ts - min_ts > (2 * 10**9):
      return

    ask0_1 = self._book_1.ask0().price
    bid0_1 = self._book_1.bid0().price
    askt_1 = self._book_1.get_notional_asks_by_qty(self._true_book_qty_1 / 100.)[1][-1][0]
    bidt_1 = self._book_1.get_notional_bids_by_qty(self._true_book_qty_1 / 100.)[1][-1][0]

    # Bitfinex -> price
    try:
      ask0_2 = self._book_2.ask0().price
      bid0_2 = self._book_2.bid0().price
    except IndexError:  # Broken book
      return
    askt_2 = self._book_2.get_notional_asks_by_amt(self._true_book_qty_2)[1][-1][0]
    bidt_2 = self._book_2.get_notional_bids_by_amt(self._true_book_qty_2)[1][-1][0]

    ask_diff = askt_1 - askt_2
    bid_diff = bidt_1 - bidt_2

    self._ask_basis_ma.update(book_ts, ask_diff)
    self._bid_basis_ma.update(book_ts, bid_diff)

    if len(self._ask_basis_ma) < 2000 or len(self._bid_basis_ma) < 2000:
      return

    askt_1_proj_2 = askt_1 - float(self._ask_basis_ma.average)
    bidt_1_proj_2 = bidt_1 - float(self._bid_basis_ma.average)
    # mid_price_2 = norm_float((ask0_2 + bid0_2) / 2., self._price_precision_2)

    askp_2 = askt_1_proj_2 * (1 + self._edge)
    askp_2 = ceil_price(askp_2, self._tick_size_2)

    bidp_2 = bidt_1_proj_2 * (1 - self._edge)
    bidp_2 = floor_price(bidp_2, self._tick_size_2)

    if book_ts < self._start_ts:
      return

    if book_ts - self._last_fire_ts < 1000 * 10**6:  # 1 seoncd
      return

    if bidt_2 > askp_2 and self._pos > -self._stack_2:
      self._pos -= 1
      self._pnl += (self._lot_size_2 * ((askp_2 * (1. - self._fee_2)) - (askt_1 *
                                                                         (1. + self._fee_1))))
      self._num_trade += 1
      self._last_fire_ts = book_ts
      print('%s %d %.8f' % (to_datetime(book_ts), self._pos, self._pnl))
    elif askt_2 < bidp_2 and self._pos < self._stack_2:
      self._pos += 1
      self._pnl += (self._lot_size_2 * ((bidt_1 * (1. - self._fee_1)) - (bidp_2 *
                                                                         (1. + self._fee_2))))
      self._num_trade += 1
      self._last_fire_ts = book_ts
      print('%s %d %.8f' % (to_datetime(book_ts), self._pos, self._pnl))

    self._min_pnl = min(self._pnl, self._min_pnl)
    self._max_pnl = max(self._pnl, self._max_pnl)

    return
    print('%d,"%s",%.8f,%.8f,%.8f,%.8f,%.8f,%.8f,%.8f,%.8f,%.8f,%.8f,%d,%.8f' %
          (book_ts,
           str(to_datetime(book_ts)),
           ask0_1,
           bid0_1,
           askt_1,
           bidt_1,
           ask0_2,
           bid0_2,
           askt_2,
           bidt_2,
           self._ask_basis_ma.average,
           self._bid_basis_ma.average,
           self._pos,
           self._pnl))


def main(argv):
  product_1 = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')
  fee_2 = 7.5 / 10000.

  product_2 = BitfinexProduct.FromStr('BTC-USDT')
  fee_1 = 0. / 10000.

  minute_ts = 60 * (10**9)
  lot_size = 0.1
  '''
  params = [
      (5, 10, 10),
      (5, 15, 10),
      (5, 20, 10),
      (5, 25, 10),
      (10, 10, 10),
      (10, 15, 10),
      (10, 20, 10),
      (10, 25, 10),
      (15, 10, 10),
      (15, 15, 10),
      (15, 20, 10),
      (15, 25, 10),
      (30, 10, 10),
      (30, 15, 10),
      (30, 20, 10),
      (30, 25, 10),
      (60, 20, 5),
      (60, 20, 7),
      (60, 20, 10),
      (60, 20, 15),
      (60, 20, 20),
      (60, 20, 25),
      (60, 20, 30)
  ]
  '''

  start_time = datetime.datetime(2018, 2, 26, 9, 0, 0)
  end_time = start_time + datetime.timedelta(hours=24)

  features = {(ma_window, edge, stack): Feature(product_1,
                                                product_2,
                                                fee_1,
                                                fee_2,
                                                edge / 10000.,
                                                ma_window * minute_ts,
                                                lot_size,
                                                stack,
                                                start_ts=start_time + datetime.timedelta(minutes=0))
              for ma_window in (10, 20, 30, 60, 90, 120) for edge in (5, 10, 15, 20, 25, 30)
              for stack in (5, 10, 15, 20, 25)}

  def on_book_reset(book_builder_name, book_builder):
    for feature in features.values():
      feature.on_book_reset(book_builder_name, book_builder)

  run_from_archive([product_1, product_2],
                   on_book_reset,
                   start_time,
                   end_time,
                   machine=['feed-3.eu-west-1'])

  for param in sorted(features.keys()):
    feature = features[param]
    logging.info('%12s pnl=%5.2f (%.2f,%.2f) norm_pnl=%5.2f %d' % (str(param),
                                                                   feature._pnl,
                                                                   feature._min_pnl,
                                                                   feature._max_pnl,
                                                                   feature._pnl / float(param[2]),
                                                                   feature._num_trade))
  return 0


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  try:
    sys.exit(main(sys.argv))
  except KeyboardInterrupt:
    sys.exit(1)
